package io.sqrtqiezi.spark.graphx

import org.apache.spark.graphx.{Edge, Graph, VertexId}
import org.apache.spark.{SparkConf, SparkContext}

case class User(name: String, age: Int, inDeg: Int, outDeg: Int)

object GraphXDemo {
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf()
      .setAppName(this.getClass.getCanonicalName)
      .setMaster( "local[*]")
    val sc = new SparkContext(conf)

    // 定义顶点 (Long, info)
    val vertexArray: Array[(VertexId, (String, Int))] = Array(
      (1L, ("Alice", 28)),
      (2L, ("Bob", 27)),
      (3L, ("Charlie", 65)),
      (4L, ("David", 42)),
      (5L, ("Ed", 55)),
      (6L, ("Fran", 50))
    )

    // 定义边 (Long, Long, attr)
    val edgeArray: Array[Edge[Int]] = Array(
      Edge(2L, 1L, 7),
      Edge(2L, 4L, 2),
      Edge(3L, 2L, 4),
      Edge(3L, 6L, 3),
      Edge(4L, 1L, 1),
      Edge(5L, 2L, 2),
      Edge(5L, 3L, 8),
      Edge(5L, 6L, 3)
    )

    // 构造 RDD
    val vertexRDD = sc.makeRDD(vertexArray)
    val edgeRDD = sc.makeRDD(edgeArray)

    // 构造图
    val graph: Graph[(String, Int), Int] = Graph(vertexRDD, edgeRDD)

    // 找出图中年龄大于 30 的顶点
    graph.vertices
      .filter{ case (_, (_, age)) => age > 30 }
      .foreach(println)

    println("------------ 分隔符 -----------------")

    // 找出图中属性大于 5 的边
    graph.edges
      .filter(_.attr > 5)
      .foreach(println)

    println("------------ 分隔符 -----------------")

    // 找出边属性大于 5 的 triplets
    graph.triplets
      .filter(_.attr > 5)
      .foreach(println)

    // degrees 操作
    // 找出图中的最大出度、入度、度数
    println("------------ 出度 -----------------")
    graph.outDegrees.foreach(println)
    val outDegrees = graph.outDegrees
      .reduce { (x, y) => if (x._2 > y._2) x else y }
    println(s"outDegrees = $outDegrees")

    println("------------ 入度 -----------------")
    graph.inDegrees.foreach(println)
    val inDegrees = graph.inDegrees
      .reduce { (x, y) => if (x._2 > y._2) x else y }
    println(s"inDegrees = $inDegrees")

    println("------------ 度数 -----------------")
    graph.degrees.foreach(println)
    val degrees = graph.degrees
      .reduce { (x, y) => if (x._2 > y._2) x else y }
    println(s"degrees = $degrees")

    println("------------ 顶点转换，所有人年龄加10岁 -----------------")
    graph.mapVertices{ case (id, (name, age)) => (id, (name, age + 10))}
      .vertices
      .foreach(println)

    println("------------ 边的转换，所有边的属性*2 -----------------")
    graph.mapEdges(_.attr*2)
      .edges
      .foreach(println)

    println("------------ 结构操作，顶点年龄 > 30 的子图 -----------------")
    val subGraph = graph.subgraph(vpred = (_, vd) => vd._2 >= 30)

    subGraph.edges.foreach(println)
    subGraph.vertices.foreach(println)

    println("------------ 连接操作，创建一个新图，顶点的数据类型为 User，并从 graph 做类型转换")
    val initialUserGraph = graph.mapVertices { case (_, (name, age)) => User(name, age, 0, 0) }

    val userGraph = initialUserGraph.outerJoinVertices(initialUserGraph.inDegrees) {
      case (_, u, inDegOpt) => User(u.name, u.age, inDegOpt.getOrElse(0), u.outDeg)
    }.outerJoinVertices(initialUserGraph.outDegrees) {
      case (_, u, outDegOpt) => User(u.name, u.age, u.inDeg, outDegOpt.getOrElse(0))
    }

    userGraph.vertices.foreach(println)

    println("----- 找到出度等于入度的人员")
    userGraph.vertices.filter{ case (_, u) => u.inDeg == u.outDeg }
      .foreach(println)

    println("------ 聚合操作，找到5到各顶点的最短距离")
    val sourceId: VertexId = 5L

    val initialGraph: Graph[Double, Int] = graph
      .mapVertices((id, _) => if (id == sourceId) 0.0 else Double.PositiveInfinity)

    val sssp = initialGraph.pregel(Double.PositiveInfinity)(
      // 两个消息来的时候，取他们中路径的最小值
      (id, dist, newDist) => math.min(dist, newDist),

      // Send Message 函数
      // 如果 triplet.srcAttr + triplet.attr < triplet.dstAttr，则发送消息到目标的顶点
      triplet => {
        if (triplet.srcAttr + triplet.attr < triplet.dstAttr) {
          Iterator((triplet.dstId, triplet.srcAttr + triplet.attr))
        } else {
          Iterator.empty
        }
      },

      // mergeMsg
      (a, b) => math.min(a, b)
    )

    println(sssp.vertices.collect.mkString("\n"))

    sc.stop()
  }
}
